#ifndef _PHOTON_MAP_H_
#define _PHOTON_MAP_H_

//
#include "spectrum.h"
#include "geometry.h"
#include "kdtree.h"
#include "sampling.h"
#include "random_number_generator.h"
#include "primitive_funcs.h"
#include "integrator_funcs.h"

#include "hybrid_lookup_accelerator.h"

struct small_point3f_t{
	float x,y,z;
	
	float& operator[](int i)
	{
		assert(i>=0 && i <=3);
		if(i == 0)return x;
		else if(i == 1)return y;
		else return z;
	}
	float operator[](int i)const
	{
		assert(i>=0 && i <=3);
		if(i == 0)return x;
		else if(i == 1)return y;
		else return z;
	}
	 operator point3f_t()
	{
		point3f_t t;
		vinit(t,x,y,z);
		return t;
	}
};
struct small_normal3f_t{
	float x,y,z;
	operator normal3f_t()
	{
		normal3f_t v;
		vinit(v,x,y,z);
		return v;
	}
};
struct small_vector3f_t{
	float x,y,z;
	operator vector3f_t()
	{
		vector3f_t v;
		vinit(v,x,y,z);
		return v;
	}
};

struct small_spectrum_t{
	float x,y,z;
	operator spectrum_t()
	{
		spectrum_t t;
		vinit(t,x,y,z);
		return t;
	}
};
template<typename T,typename U>
inline float distance_squared(const T& p0, const U& p1)
{
	small_vector3f_t v;
	vsub(v,p0,p1);
	return vdot(v,v);
}
INLINE void bbox_union_with_point(bbox_t* bbox,const small_point3f_t* p)
{
	bbox->pmin.x = min(bbox->pmin.x, (*p).x);
    bbox->pmin.y = min(bbox->pmin.y, (*p).y);
    bbox->pmin.z = min(bbox->pmin.z, (*p).z);
    bbox->pmax.x = max(bbox->pmax.x, (*p).x);
    bbox->pmax.y = max(bbox->pmax.y, (*p).y);
    bbox->pmax.z = max(bbox->pmax.z, (*p).z);
}
typedef struct{
	small_point3f_t p;
	small_spectrum_t alpha;
	small_vector3f_t wi;
} photon_t;

INLINE void photon_init(photon_t* photon,const point3f_t* p,const spectrum_t* wt,const vector3f_t* w)
{
	vassign(photon->p , (*p));
	vassign(photon->alpha,(*wt));
	vassign(photon->wi , (*w));;
}

typedef struct {
	small_point3f_t p;
	small_normal3f_t n;
	small_spectrum_t lo;
}radiance_photon_t ;

INLINE void radiance_photon_init(radiance_photon_t* photon,const point3f_t* p,const normal3f_t* n)
{
	vassign(photon->p , (*p));
	vassign(photon->n , (*n));
	vclr(photon->lo);
}

typedef struct{
	GLOBAL photon_t* photon;
	float distance_squared;
}
close_photon_t;

INLINE void close_photon_init(LOCAL close_photon_t* close_photon,GLOBAL photon_t* photon,float dist_sqr)
{
	close_photon->photon = photon;
	close_photon->distance_squared = dist_sqr;
}
INLINE bool close_photon_less_than(close_photon_t a,close_photon_t b)
{
	return a.distance_squared == b.distance_squared? (a.photon < b.photon) :
		a.distance_squared < b.distance_squared;
}
typedef struct  
{
	LOCAL close_photon_t *photons;
	unsigned n_lookup;
	unsigned found_photons;
}photon_process_data_t;

INLINE void photon_process_data_init(photon_process_data_t* photon_process_data,
						 unsigned mp
						 )
{
	photon_process_data->photons = 0;
	photon_process_data->n_lookup = mp;
	photon_process_data->found_photons = 0;
}
INLINE void photon_process(photon_process_data_t* data,GLOBAL photon_t* photon,
					float dist2,float* max_dist_sqred)
{
	if (data->found_photons < data->n_lookup) {
		// Add photon to unordered array of photons
		close_photon_init(&data->photons[data->found_photons++] ,photon, dist2);
		if (data->found_photons == data->n_lookup) {
			//c_make_heap(close_photon_t,&data->photons[0], &data->photons[data->n_lookup],close_photon_less_than);
			std::make_heap(&data->photons[0], &data->photons[data->n_lookup],close_photon_less_than);
			*max_dist_sqred = data->photons[0].distance_squared;
		}
	}
	else {
		// Remove most distant photon from heap and add new photon
		LOCAL close_photon_t *begin = data->photons;
		LOCAL close_photon_t *end = data->photons + data->n_lookup;
		//c_pop_heap(close_photon_t,begin, end,close_photon_less_than);
		std::pop_heap(begin,end,close_photon_less_than);
		close_photon_init(&data->photons[data->n_lookup-1] ,photon, dist2);
		//c_push_heap(close_photon_t,begin, end,close_photon_less_than);
		std::push_heap(begin, end,close_photon_less_than);
		*max_dist_sqred = data->photons[0].distance_squared;
	}
}

typedef struct{
	const point3f_t *p;
	const normal3f_t *n;
	GLOBAL radiance_photon_t *photon;
} radiance_photon_process_data_t;
INLINE void radiance_photon_process_data_init(radiance_photon_process_data_t *data,const point3f_t* pp,const normal3f_t *nn)
{
	data->p = pp;
	data->n = nn;
	data->photon = 0;
}

INLINE void radiance_photon_process(radiance_photon_process_data_t *data,GLOBAL radiance_photon_t *rp,
							 float dist_sqr,float *max_dist)
{
	if(vdot(rp->n,*data->n)>0)
	{
		data->photon = rp;
		*max_dist = dist_sqr;
	}
}



typedef kd_tree_t(photon_t) photon_kd_tree_t;
typedef kd_tree_t(radiance_photon_t) radiance_photon_kd_tree_t;

typedef struct {
	// ExPhotonIntegrator Private Data
    unsigned n_caustic_photons, n_indirect_photons;
	unsigned  n_lookup;
	unsigned int max_specular_depth;
	float max_dist_squared, rr_threshold;
	// Declare sample parameters for light source sampling
	unsigned int n_caustic_paths, n_indirect_paths;
	HybridLookupAccelerator<photon_t>* indirect_map_lookup;
	HybridLookupAccelerator<photon_t>* caustic_map_lookup;
	HybridLookupAccelerator<radiance_photon_t>* radiance_map_lookup;

	unsigned total_photons;
	//
    unsigned int final_gather;
	float cos_gather_angle;
	unsigned int gather_samples;
	//ppm data
	unsigned progressive_iteration;
	float alpha;
}
photon_map_t;


INLINE float photon_map_kernel(GLOBAL photon_t *photon, const point3f_t *p,
		float md2) {
	float s = (1.f - distance_squared(photon->p, *p) / md2);
	return 3.f / (md2 * FLOAT_PI) * s * s;
}

#define MAX_CLOSE_PHOTON_LOOKUP 512

INLINE float photon_map_get_radius(const photon_map_t* photon_map,float reference_radius_squared)
{

	float radius_squared = reference_radius_squared;
	//return max(0.0001f,radius_squared);
	return radius_squared;
}

INLINE void photon_map_lphoton(const photon_map_t* photon_map,
							   const HybridLookupAccelerator<photon_t> *map,
							   int n_paths,
							   int n_lookup,
							   bsdf_t *bsdf,
							   const intersection_t* isect,
							   const vector3f_t *wo,
							   float max_dist_sqr,
							   cl_scene_info_t scene_info,
							   Seed *seed,
							   spectrum_t *f,
							   LOCAL close_photon_t *close_photon_data_store
							   )
{
	if(!map || map->Empty()) return;
	vclr(*f);
	int non_specular = BSDF_REFLECTION | BSDF_TRANSMISSION |  BSDF_DIFFUSE | BSDF_GLOSSY;
	if (bsdf_num_components(bsdf,non_specular) == 0)
	{
		return;
	}
	photon_process_data_t photon_map_data;
	photon_process_data_init(&photon_map_data,photon_map->n_lookup);
	photon_map_data.photons = close_photon_data_store;
	//kd_tree_lookup(*map,isect->dg.p,&photon_map_data,photon_process,max_dist_sqr);
	float radius_i_squared = photon_map_get_radius(photon_map,max_dist_sqr);
	map->Lookup(isect->dg.p,photon_map_data,photon_process,radius_i_squared);
	//estimate reflected light from photons
	LOCAL close_photon_t *photons = photon_map_data.photons;
	unsigned int n_found = photon_map_data.found_photons;
	normal3f_t n_f;
	if(vdot(*wo,bsdf->dg_shading.nn) < 0)
	{
		vneg(n_f,bsdf->dg_shading.nn);
	}
	else
		vassign(n_f,bsdf->dg_shading.nn);
	//float radius_i_squared = n_found >= photon_map->n_lookup ? 
		//photon_map_get_radius(photon_map,max_dist_sqr):max_dist_sqr;
	if (bsdf_num_components(bsdf,BSDF_REFLECTION|BSDF_TRANSMISSION|BSDF_GLOSSY)> 0)
	{
		spectrum_t L;
		//compute exitant radiance from photons for glossy
		for(unsigned int i = 0;i < n_found; ++i)
		{
			//if(photons[i].distance_squared > radius_i_squared)
			//	continue;
			GLOBAL photon_t *p = photons[i].photon;
			BxDFType flag = vdot(n_f,p->wi) > 0 ?BSDF_ALL_REFLECTION:BSDF_ALL_TRANSMISSION;
			float k = photon_map_kernel(p,&isect->dg.p, radius_i_squared);
			vector3f_t wi_tmp ;
			vassign(wi_tmp,p->wi);
			bsdf_f(bsdf,wo,&wi_tmp,flag,&L);
			vmul(L,L,p->alpha);
			vsmul(L,(k/n_paths),L);
			vadd(*f,*f,L);
		}
	}
	else
	{
		spectrum_t Lr,Lt;
		vclr(Lr);vclr(Lt);
		spectrum_t L;
		for (unsigned int i = 0;i < n_found; ++i)
		{

			//if(photons[i].distance_squared > radius_i_squared)
			//	continue;
			float k = photon_map_kernel(photons[i].photon,&isect->dg.p,
				radius_i_squared);
			if(vdot(n_f,photons[i].photon->wi) > 0.f)
			{
				vsmul(L,(k/n_paths),photons[i].photon->alpha);
				vadd(Lr,Lr,L);
			}
			else
			{
				vsmul(L,(k/n_paths),photons[i].photon->alpha);
				vadd(Lt,Lt,L);
			}
		}
		bsdf_rho_hd(bsdf,wo,seed,BSDF_ALL_REFLECTION,&L);
		vmul(Lr,Lr,L);
		vsmul(Lr,INV_PI,Lr);
		bsdf_rho_hd(bsdf,wo,seed,BSDF_ALL_TRANSMISSION,&L);
		vmul(Lt,Lt,L);
		vsmul(Lt,INV_PI,Lt);
		vadd(*f,*f,Lr);
		vadd(*f,*f,Lt);
	}
}
#define MAX_RAY_DEPTH 4
//todo optimization
INLINE void photon_map_final_gather(const photon_map_t* photon_map,cl_scene_info_t scene_info,
									Seed* seed,
									const point3f_t p,
									const normal3f_t n,
									const vector3f_t wo,
									bsdf_t *bsdf,
									spectrum_t *l,
									LOCAL close_photon_t *photon_buffer,
									LOCAL vector3f_t *photon_dirs
									)
{
	int non_specular = BSDF_REFLECTION|BSDF_TRANSMISSION|BSDF_DIFFUSE|BSDF_GLOSSY;
	if (bsdf_num_components(bsdf,non_specular)>0)
	{
#define n_indir_sample_photons 50

		 photon_process_data_t proc_data;
		photon_process_data_init(&proc_data,
			n_indir_sample_photons);
		proc_data.photons = photon_buffer;
		float search_dist2 = photon_map->max_dist_squared;
		int debug_count = 0;
		while (proc_data.found_photons < n_indir_sample_photons)//dead loop
		{
			float md2 = search_dist2;
			proc_data.found_photons = 0;
			photon_map->indirect_map_lookup->Lookup(p,proc_data,photon_process,md2);
			search_dist2 *= 2.f;
			debug_count ++;
			if (debug_count > 32)
			{
				printf("debug_count dead loop when search photons nearby total %d \n", photon_map->indirect_map_lookup->Size());
				return;
			}
		}
		for (unsigned i = 0;i < n_indir_sample_photons; ++i)
		{
			vassign(photon_dirs[i],proc_data.photons[i].photon->wi);
		}
		spectrum_t Li;vclr(Li);
		for(unsigned int i = 0; i < photon_map->gather_samples; ++i)
		{
			vector3f_t wi;float pdf;
			spectrum_t fr;
			float u1 = random_float(seed);
			float u2 = random_float(seed);
			float u3 = random_float(seed);
			BxDFType sampled_type;
			bsdf_sample_f(bsdf,&wo,&wi,u1,u2,u3,&pdf,(BxDFType)(BSDF_ALL&~BSDF_SPECULAR),&sampled_type,&fr);
			if (pdf == 0.f || color_is_black(fr))
			{
				continue;
			}
			ray_t bounce_ray;
			rinit(bounce_ray,p,wi);
			intersection_t gather_isect;
			if (intersect(scene_info.accelerator_data,scene_info.shape_data,scene_info.primitives,
				scene_info.primitive_count,&bounce_ray,&gather_isect))
			{
				spectrum_t L_indir;vclr(L_indir);
				normal3f_t normal_gather = gather_isect.dg.nn;
				if(vdot(normal_gather,bounce_ray.d)>0)vneg(normal_gather,normal_gather);
				radiance_photon_process_data_t  radiance_proc_data;
				radiance_photon_process_data_init(&radiance_proc_data,
					&gather_isect.dg.p,&normal_gather);
				float md2 = FLT_MAX;
				photon_map->radiance_map_lookup->Lookup(
					gather_isect.dg.p,radiance_proc_data,
					radiance_photon_process,
					md2
					);
				if(radiance_proc_data.photon) L_indir = radiance_proc_data.photon->lo;
				spectrum_t trans;
				scene_tranmittance(scene_info,&bounce_ray,&trans);
				vmul(L_indir,L_indir,trans);

				float photon_pdf = 0.f;
				float cone_pdf = uniform_cone_pdf(photon_map->cos_gather_angle);
				for (int j = 0; j < n_indir_sample_photons;++j)
				{
					if(vdot(photon_dirs[j],wi) > 0.999f * photon_map->cos_gather_angle)photon_pdf += cone_pdf;
				}
				photon_pdf /= n_indir_sample_photons;
				float wt = power_heuristic(photon_map->gather_samples,pdf,photon_map->gather_samples,photon_pdf);
				vsmul(L_indir,(wt*fabs(vdot(wi,n))/pdf),L_indir);
				vmul(fr,fr,L_indir);
				vadd(Li,fr,Li);

			}
		}
		vsmul(Li,1.f/photon_map->gather_samples,Li);
		vadd(*l,*l,Li);
		vclr(Li);
		for (unsigned int i = 0;i < photon_map->gather_samples;++i)
		{
			float u1 = random_float(seed);
			float u2 = random_float(seed);
			float u3 = random_float(seed);
			int photon_num = 
				min((int)n_indir_sample_photons -1,
				(int)(u1 * n_indir_sample_photons));
			vector3f_t vx,vy;
			vector3f_t photon_dir = photon_dirs[photon_num];
			coordinate_system(&photon_dir,&vx,&vy);
			vector3f_t wi;
			uniform_sample_cone(u2,u3,
				photon_map->cos_gather_angle,
				&vx,&vy,&photon_dir,
				&wi);
			spectrum_t fr;
			bsdf_f(bsdf,&wo,&wi,BSDF_ALL,&fr);
			if (color_is_black(fr))
				continue;
			ray_t bounce_ray;
			rinit(bounce_ray,p,wi);
			intersection_t gather_isect;
			if(intersect(scene_info.accelerator_data,scene_info.shape_data,scene_info.primitives,
				scene_info.primitive_count,&bounce_ray,&gather_isect))
			{
				spectrum_t L_indir;vclr(L_indir);
				normal3f_t nn;vassign(nn,gather_isect.dg.nn);
				if(vdot(nn,bounce_ray.d) > 0.f) vneg(nn,nn);
				radiance_photon_process_data_t  radiance_proc_data;
				radiance_photon_process_data_init(&radiance_proc_data,
					&gather_isect.dg.p,&nn);
				float md2 = FLT_MAX;
				photon_map->radiance_map_lookup->Lookup(
					gather_isect.dg.p,radiance_proc_data,
					radiance_photon_process,
					md2
					);
				if(radiance_proc_data.photon)
					vassign(L_indir,radiance_proc_data.photon->lo );
				spectrum_t trans;
				scene_tranmittance(scene_info,&bounce_ray,&trans);
				vmul(L_indir,L_indir,trans);


				float photon_pdf = 0.f;
				float cone_pdf = uniform_cone_pdf(
					photon_map->cos_gather_angle);
				for (unsigned j = 0; j < n_indir_sample_photons;++j)
				{
					if(vdot(photon_dirs[j],wi) > 0.999f * photon_map->cos_gather_angle)
						photon_pdf += cone_pdf;
				}
				photon_pdf /= n_indir_sample_photons;

				float bsdfPdf = bsdf_pdf(bsdf,&wo,&wi,BSDF_ALL);
				float wt = power_heuristic(photon_map->gather_samples,
					photon_pdf,
					photon_map->gather_samples,
					bsdfPdf);
				vsmul(L_indir,(wt*fabs(vdot(wi,n))/photon_pdf),L_indir);
				vmul(fr,fr,L_indir);
				vadd(Li,fr,Li);
			}
		}
		vsmul(Li,1.f/photon_map->gather_samples,Li);
		vadd(*l,*l,Li);
	}
}
/*
INLINE void direct_lighting(spectrum_t *li,cl_scene_info_t scene_info, const ray_differential_t *cur_ray,Seed* seed,const intersection_t *isect)
{
	vector3f_t wo;
	vneg(wo,cur_ray->d);
	//compute emitted light if ray hit an area light
	spectrum_t l;
	spectrum_t li_val;
	bsdf_t bsdf;
	vclr(li_val);
	vclr(l);
	intersection_le(&isect,scene_info,&wo,&l);
	vadd(li_val,li_val,l);
	intersection_get_bsdf(&isect,scene_info,
		cur_ray,&bsdf);
	point3f_t p = bsdf.dg_shading.p;
	normal3f_t n = bsdf.dg_shading.nn;
	uniform_sample_all_lights(scene_info,
		seed,&p,&n,&wo,&bsdf,&l);
	vadd(li_val,li_val,l);
	*li = li_val;
}
INLINE void photon_map_lphoton_lookup(photon_map_t *photon_map,const photon_kd_tree_t* map,unsigned int n_paths,  
const intersection_t* isect,const ray_differential_t *cur_ray, cl_scene_info_t scene_info,Seed *seed,spectrum_t *l)
{
	vector3f_t wo;
	spectrum_t l_tmp;
	vneg(wo,cur_ray->d);
	bsdf_t bsdf;
	intersection_get_bsdf(&isect,scene_info,
		cur_ray,&bsdf);

	photon_map_lphoton(photon_map,
		map,n_paths,
		photon_map->n_lookup,&bsdf,isect,&wo,photon_map->max_dist_squared
		,scene_info,seed,&l_tmp);
	vadd(*l,l_tmp,*l);
}
typedef struct{
	ray_t ray_stack[MAX_RAY_DEPTH];
    spectrum_t passthrough[MAX_RAY_DEPTH];
	bool left_stack[MAX_RAY_DEPTH];
typedef enum _VisitType {pre,in,post} VisitType;
	VisitType visit;
	int ray_stack_top;
}intersect_ray_stack_t;

typedef struct{
	intersection_t isect;
	spectrum_t passthrough;
}
intersect_ray_output_t
;
INLINE void ray_buffer_intersect(intersect_ray_output_t *ray_out,const intersect_ray_stack_t *ray_stack,
	cl_scene_info_t scene_info)
{
	ray_t cur_ray;
	intersection_t isect;
	while(ray_stack->ray_stack_top > 0)
	{
		rassign(cur_ray,ray_stack->ray_stack[ray_stack->ray_stack_top - 1]);
		if(ray_stack->visit != post && intersect(scene_info.accelerator_data,scene_info.shape_data,scene_info.primitives,
			scene_info.primitive_count,&cur_ray,&isect)
			)
		{
			if(ray_stack->visit == pre)
			{
				//todo: store output data
			}
			if(ray_stack->ray_stack_top < photon_map->max_specular_depth && ray_stack->visit != post)
			{
                if(ray_stack->visit == pre)
                {
                    specular_reflect(&cur_ray,bsdf_stack+ray_stack_top-1,seed,&isect,scene_info,passthrough+ray_stack_top,
                            ray_stack+ray_stack_top);
                    if(color_is_black(passthrough[ray_stack_top]))
                    {
                        //go on try transmission rays
                        specular_transmit(&cur_ray,bsdf_stack + ray_stack_top -1,seed,&isect,scene_info,passthrough+ray_stack_top,
                                ray_stack+ray_stack_top);
                        if(color_is_black(passthrough[ray_stack_top]))
                        {
                            visit = post;
                        }
                        else
                        {
							vmul(*(passthrough+ray_stack_top),*(passthrough+ray_stack_top),*(passthrough+ray_stack_top-1));
							left_stack[ray_stack_top] = false;
                            visit = pre;
                            ray_stack_top++;
                        }
                    }
                    else{
                        //multiply passthrough
						vmul(*(passthrough+ray_stack_top),*(passthrough+ray_stack_top),*(passthrough+ray_stack_top-1));
						left_stack[ray_stack_top] = true;
                        visit = pre;
                        ray_stack_top++;
                    }
                }
                else if(visit == in)
                {
                    //go on try transmission rays
                    specular_transmit(&cur_ray,bsdf_stack + ray_stack_top - 1,seed,&isect,scene_info,passthrough+ray_stack_top,
                            ray_stack+ray_stack_top);
                    if(color_is_black(passthrough[ray_stack_top]))
                    {
                        visit = post;
                    }
                    else
                    {
						vmul(*(passthrough+ray_stack_top),*(passthrough+ray_stack_top),*(passthrough+ray_stack_top-1));
						left_stack[ray_stack_top] = false;
                        visit = pre;
                        ray_stack_top++;
                    }
                }
				return;
		}
		else
		{
			//handle no intersection
			if(ray_stack->visit == in)ray_stack->visit = post;
			else if(ray_stack->visit == post && ray_stack->left_stack[ray_stack->ray_stack_top-1])
				ray_stack->visit = in;
			else if(ray_stack->visit == pre)
				if(ray_stack->left_stack[ray_stack->ray_stack_top-1])ray_stack->visit = in;
				else ray_stack->visit = post;

			ray_stack->ray_stack_top--;

			//ray_stack_top--;
			//visit = pre;
		}
	}
}
*/

#ifndef CL_KERNEL
#include "random_number_generator_mt19937.h"
#include "scene_data.h"


void estimate_e(photon_map_t* photon_map,HybridLookupAccelerator<photon_t>* map,int count,
				const point3f_t& p, const normal3f_t& n,
				spectrum_t* ret);
bool unsuccessful(int needed, int found, int shot);

 void photon_map_destroy(photon_map_t* photon_map);//free up kd-tree memory
#endif


#endif
